
# coding: utf-8

# In[1]:


import tensorflow as tf
import sys
sys.path.append('../')
import train.config as config


# In[2]:


class FcnDetector:
    '''识别单张图片'''
    def __init__(self,net_factory,model_path):
        model = net_factory()
        ckpt = tf.train.Checkpoint(model=model)
        ckpt_manager = tf.train.CheckpointManager(ckpt, model_path, max_to_keep=5)
        ckpt.restore(ckpt_manager.latest_checkpoint)
        self.model = model

    def predict(self,databatch):
        data = tf.convert_to_tensor(databatch, dtype=tf.float32)
        data = tf.expand_dims(data, 0)
        cls_prob, bbox_pred, _ = self.model(data)
        
        return cls_prob.numpy(),bbox_pred.numpy()

        

